﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using Microsoft.WindowsAzure;

namespace EncryptDecrypt
{
    internal class AzureTableCryptoKeyStore : IDisposable
    {
        private Dictionary<int, SymmetricAlgorithm> keyCache = new Dictionary<int, SymmetricAlgorithm>();

        internal AzureTableCryptoKeyStore(CloudStorageAccount acct)
        {
            SymmetricKeyDataServiceContext keyTable = new SymmetricKeyDataServiceContext(acct);
            List<SymmetricKey> allKeys = keyTable.SymmetricKeys.ToList();


            foreach (var key in allKeys)
            {
                X509Certificate2 certificate = GetCertificate(key.CertificateThumbprint);
                RSACryptoServiceProvider RSA = (RSACryptoServiceProvider)certificate.PrivateKey;
                byte[] symmetricCryptoKey = RSA.Decrypt(key.Key, true);

                AesManaged algorithm = new AesManaged();
                algorithm.IV = key.iv;
                algorithm.Key = symmetricCryptoKey;
                keyCache[key.Version] = algorithm;
            }
        }

        internal ICryptoTransform GetDecryptor(int version)
        {
            return GetAlgorithm(version).CreateDecryptor();
        }

        internal ICryptoTransform GetEncryptor(int version)
        {
            return GetAlgorithm(version).CreateEncryptor();
        }

        private SymmetricAlgorithm GetAlgorithm(int version)
        {
            SymmetricAlgorithm algo;
            if (!keyCache.TryGetValue(version, out algo))
            {
                throw new InvalidOperationException("Could not find Azure Table Encryption version " + version);
            }
            return algo;
        }

        public static X509Certificate2 GetCertificate(string thumbprint, StoreLocation storeLocation = StoreLocation.LocalMachine, StoreName storeName = StoreName.My)
        {
            if (string.IsNullOrEmpty(thumbprint))
            {
                throw new ArgumentNullException("thumbprint");
            }

            thumbprint = thumbprint.Replace(" ", "").ToUpperInvariant();

            X509Store store = new X509Store(storeName, storeLocation);

            try
            {
                store.Open(OpenFlags.ReadOnly);

                foreach (var cert in store.Certificates)
                {
                    if (cert.HasPrivateKey && String.Compare(cert.Thumbprint, thumbprint) == 0)
                    {
                        return cert;
                    }
                }
            }
            finally
            {
                store.Close();
            }

            return null;
        }

        public void Dispose()
        {
            Dictionary<int, SymmetricAlgorithm> cache = keyCache;
            keyCache = null;

            foreach (var algo in cache.Values)
            {
                algo.Clear();
                algo.Dispose();
            }
        }
    }
}
